#include "ros/ros.h"
#include <geometry_msgs/Twist.h>         // 指线速度角速度的消息类型
#include <tf/transform_listener.h>       // 用于坐标转换
#include <nav_msgs/Path.h>                  // 显示历史轨迹和路径
#include <std_msgs/Float32.h>
#include <std_msgs/Int32.h>
#include <nav_msgs/Odometry.h>           // 用于读取里程计信息
#include <sensor_msgs/LaserScan.h>       // lidar
#include <robust_formation/RobustleaderMpc.h>                 // MPC解优化
#include <Eigen/Core>                    // 矩阵运算
#include <Eigen/QR>

// inlcude iostream and string libraries
#include <iostream>
#include <map>      // 包含了key-value键值对映射的排序分配容器
#include <math.h>
#include <string>
#include <vector>

using namespace std;
using namespace Eigen;

constexpr double pi() { return M_PI; }
double deg2rad(double x) { return x * pi() / 180; }
double rad2deg(double x) { return x * 180 / pi(); }

/********************/
/* CLASS DEFINITION */
/********************/
class RobustMPCNode
{
    public:
        RobustMPCNode();  // 构造函数
        ~RobustMPCNode(); // 析构函数
        int get_thread_numbers();   // 获取线程数
    
    private:
        ros::NodeHandle _nh;
        ros::Subscriber _sub_odom, _sub_distance, _sub_cmd, _sub_cleantime;
        ros::Publisher _pub_totalcost, _pub_twist, _pub_mpctraj, _pub_mpcdif;
        ros::Timer _timer1;
        nav_msgs::Odometry _odom;
        nav_msgs::Path _mpc_traj;
        geometry_msgs::Twist _twist_msg, _sub_twist_cmd;

        string _map_frame, _odom_frame, _car_frame;
        string _leader_frame;

    { //用于可视化操作
        // ✅ 新增：可视化发布者
        ros::Publisher _pub_debug_info;          // 调试信息
        ros::Publisher _pub_mpc_state;           // MPC状态可视化
        ros::Publisher _pub_obstacle_marker;     // 障碍物标记
        ros::Publisher _pub_robot_trajectory;    // 机器人轨迹
        ros::Publisher _pub_prediction_traj;     // 预测轨迹
    
        // ✅ 新增：调试消息类型
        std_msgs::String _debug_msg;
        visualization_msgs::Marker _obstacle_marker;
        nav_msgs::Path _robot_path;
        
        // ✅ 新增：调试计数器
        int _debug_counter;
        double _last_debug_time;
        // ✅ 新增：可视化工具函数
        void publishDebugInfo(const std::string& function_name, const std::string& message, int level=0);
        void publishMPCStateVisualization(const Eigen::VectorXd& state, const Eigen::VectorXd& cmd);
        void publishObstacleMarker(double distance, double angle);
        void publishRobotTrajectory(double x, double y, double theta);
        void publishPredictionTrajectory(const std::vector<double>& mpc_x, const std::vector<double>& mpc_y);
        std::string getCurrentTimeString();
        
        // ✅ 新增：状态检查函数
        void checkMPCInputs(const Eigen::VectorXd& state, const Eigen::VectorXd& cmd);
        void validateMPCOutputs(const std::vector<double>& results);
    }
        Robust_BCI_MPC _mpc;
        // map<string, double> _mpc_params;
        // double _mpc_steps, _w_linvel, _w_linvel_d, _w_angvel, _w_angvel_d, _max_angvel, _max_speed, _bound_value;

        double _dt, _w, _speed;
        int _controller_freq, _thread_numbers;
        bool _debug_info, _delay_mode;
        float _closest_distance, _closest_degree;                               // 雷达测量障碍物的最近距离和角度
        // double _target1_x, _target1_y, _target2_x, _target2_y;                  // set target A(5, 5) B(5, -5)
        // double distanceToA, distanceToB;                                        // 任意时刻Leader与目标的距离

        void odomCB(const nav_msgs::Odometry::ConstPtr& odomMsg);               // 订阅里程计消息回调函数
        void controlLoopCB(const ros::TimerEvent&);                             // 控制回调函数
        void get_laser_callback(const sensor_msgs::LaserScan::ConstPtr& msg);   // 订阅雷达消息回调函数
        void read_cmd(const geometry_msgs::Twist::ConstPtr& cmd);               // 读取用户输入的命令
        void clean_time_CB(const std_msgs::Int32::ConstPtr& msg);               // 将当前时间归零

        // Intention tracking control
        void robust_bci_ctrl();   

        unsigned int idx; 
        double _time_begin;       // 记录开始时间
        double _cur_time;         // 当前时间
        std_msgs::Int32 cwm_flag; // 记录当前状态（到达A、B或者超时）

        geometry_msgs::Twist _mpc_dif; // 记录mpc变化情况
}; // end of class

// --------------------关键核心函数----------------------------//
RobustMPCNode::RobustMPCNode()
{
    // //Parameter for MPC solver
    // pn.param("mpc_steps", _mpc_steps, 4.0);
    // pn.param("mpc_w_angvel", _w_angvel, 5.0);
    // pn.param("mpc_w_angvel_d", _w_angvel_d, 1.0);
    // pn.param("mpc_w_linvel", _w_linvel, 100.0);
    // pn.param("mpc_w_linvel_d", _w_linvel_d, 1.0);    
    // pn.param("max_speed", _max_speed, 0.2);// unit: m/s
    // pn.param("max_angvel", _max_angvel, 0.2); // Maximal angvel radian 
    // pn.param("mpc_bound_value", _bound_value, 1.0e3); // Bound value for other variables

    
    {// Param
    ros::NodeHandle pn("~");
    pn.param("thread_numbers", _thread_numbers, 2);   //Parameters for control loop
    pn.param("debug_info", _debug_info, true);
    pn.param("delay_mode", _delay_mode, false);
    pn.param("controller_freq", _controller_freq, 4);
    _dt = double(1.0/_controller_freq);
    
    //Parameter for topics & Frame name
    pn.param<std::string>("map_frame", _map_frame, "tb3_0/odom" ); //*****for mpc, "odom"
    pn.param<std::string>("odom_frame", _odom_frame, "tb3_0/odom");
    pn.param<std::string>("car_frame", _car_frame, "tb3_0/base_footprint" );
    pn.param<std::string>("leader_frame", _leader_frame, "tb3_0" );

    // set targets
    pn.param("target1_x", _target1_x, 5.0);
    pn.param("target1_y", _target1_y, 5.0);
    pn.param("target2_x", _target2_x, 5.0);
    pn.param("target2_y", _target2_y, -5.0);

    //Display the parameters
    // cout << "\n===== Parameters =====" << endl;
    // cout << "frequency: "   << _dt << endl;
    // cout << "mpc_steps: "   << _mpc_steps << endl;
    // cout << "mpc_w_linvel: "   << _w_linvel << endl;
    // cout << "mpc_w_linvel_d: "  << _w_linvel_d << endl;    
    // cout << "mpc_w_angvel: "   << _w_angvel << endl;
    // cout << "mpc_w_angvel_d: "  << _w_angvel_d << endl;
    // cout << "mpc_max_angvel: "  << _max_angvel << endl;
    // cout << "mpc_max_speed: "  << _max_speed << endl;
    }


// ros 部分
{
    //Publishers and Subscribers
    _sub_odom   = _nh.subscribe(_leader_frame + "/odom", 1, &RobustMPCNode::odomCB, this);                   // 订阅里程计消息
    _sub_distance   = _nh.subscribe(_leader_frame + "/processed_scan", 1, &RobustMPCNode::get_laser_callback, this);   // 订阅雷达消息
    // _sub_cmd = _nh.subscribe("/cmd_udp", 1, &RobustMPCNode::read_cmd, this); // change 20210113
    _sub_cmd = _nh.subscribe("/cmd_user", 1, &RobustMPCNode::read_cmd, this); // change 20210113
    _sub_cleantime = _nh.subscribe("/publishCleanTime", 1, &RobustMPCNode::clean_time_CB, this);

    _pub_twist = _nh.advertise<geometry_msgs::Twist>(_leader_frame + "/cmd_vel", 1);    // publish speed msg
    _pub_totalcost  = _nh.advertise<std_msgs::Float32>("/total_cost", 1);               // total_cost
    _pub_mpctraj = _nh.advertise<nav_msgs::Path>("/mpc_trajectory", 1);   // MPC trajectory output  /  Prediction time domain
    _pub_mpcdif = _nh.advertise<geometry_msgs::Twist>("/mpc_dif", 1);  // 发布MPC的变化情况
    
    //Timer
    _timer1 = _nh.createTimer(ros::Duration((1.0)/_controller_freq), &RobustMPCNode::controlLoopCB, this); // 4Hz //*****mpc
}
    // //Init parameters for MPC object
    // _mpc_params["DT"] = _dt;
    // _mpc_params["STEPS"]    = _mpc_steps;
    // _mpc_params["W_LINVEL"]  = _w_linvel;
    // _mpc_params["W_DLINVEL"] = _w_linvel_d;    
    // _mpc_params["W_ANGVEL"]  = _w_angvel;
    // _mpc_params["W_DANGVEL"] = _w_angvel_d;
    // _mpc_params["MAX_ANGVEL"]   = _max_angvel;
    // _mpc_params["MAX_SPEED"]   = _max_speed;
    // _mpc_params["BOUND"]    = _bound_value;
    // _mpc.LoadParams(_mpc_params);

    idx = 0;         // 求解次数

    _w = 0.0;        // 角速度
    _speed = 0.0;    // 速度

    _closest_degree = 0;        // 初始化最近障碍物角度和距离
    _closest_distance = 3.5;

    _twist_msg = geometry_msgs::Twist();
    _mpc_traj = nav_msgs::Path();

    _time_begin = ros::Time::now().toSec();
    while (_time_begin == 0){
        _time_begin = ros::Time::now().toSec();
    }
    cout << "time_begin: " << _time_begin << endl;

    // ✅ 新增：初始化可视化发布者
    _pub_debug_info = _nh.advertise<std_msgs::String>("/debug/mpc_info", 10);
    _pub_mpc_state = _nh.advertise<visualization_msgs::Marker>("/debug/mpc_state", 10);
    _pub_obstacle_marker = _nh.advertise<visualization_msgs::Marker>("/debug/obstacle", 10);
    _pub_robot_trajectory = _nh.advertise<nav_msgs::Path>("/debug/robot_path", 10);
    _pub_prediction_traj = _nh.advertise<nav_msgs::Path>("/debug/prediction", 10);
    
    // 初始化调试变量
    _debug_counter = 0;
    _last_debug_time = ros::Time::now().toSec();
    
    ROS_INFO("✅ 可视化调试系统已初始化");

}

RobustMPCNode::~RobustMPCNode()
{
    // 析构函数  
}

// Public: return _thread_numbers
int RobustMPCNode::get_thread_numbers()
{
    return _thread_numbers;
}

// Timer: Control Loop (closed loop nonlinear MPC)
void RobustMPCNode::controlLoopCB(const ros::TimerEvent&)
{       
    geometry_msgs::Twist cmd_user = _sub_twist_cmd;
    const double autonomy = cmd_user.linear.z;
    double curent_time = ros::Time::now().toSec() - _time_begin;
    cout << "duration: " << curent_time << endl; 
    robust_bci_ctrl();
}

void RobustMPCNode::robust_bci_ctrl() {
    // ✅ 记录函数开始
    publishDebugInfo("robust_bci_ctrl", "开始执行MPC控制", 0);
        try {
        // 获取输入数据
        nav_msgs::Odometry odom = _odom;
        geometry_msgs::Twist cmd_user = _sub_twist_cmd;
        
        // ✅ 检查输入数据
        if (std::isnan(odom.pose.pose.position.x)) {
            publishDebugInfo("robust_bci_ctrl", "里程计数据包含NaN!", 2);
            return;
        }
        
        // 准备MPC输入
        VectorXd vec_cmd(2);
        vec_cmd << cmd_user.linear.x, cmd_user.angular.z;

        VectorXd state(5);
        state << 0, 0, 0, _closest_distance, _closest_degree;
        
        // ✅ 发布状态可视化
        publishMPCStateVisualization(state, vec_cmd);
        publishObstacleMarker(_closest_distance, _closest_degree);
        
        // ✅ 记录MPC求解前状态  .transpose()----->>Matrix transpose
        std::stringstream ss;
        ss << "开始MPC求解: state=[" << state.transpose() << "], cmd=[" << vec_cmd.transpose() << "]";
        publishDebugInfo("robust_bci_ctrl", ss.str(), 0);
        // 当前时间
        _cur_time = ros::Time::now().toSec() - _time_begin;
        // MPC求解
        vector<double> mpc_results = _mpc.Solve(state, vec_cmd);
        
        // ✅ 检查MPC输出
        if (mpc_results.size() < 2) {
            publishDebugInfo("robust_bci_ctrl", "MPC返回结果数量不足!", 2);
            return;
        }
        
        // ✅ 验证输出有效性
        if (std::isnan(mpc_results[0]) || std::isnan(mpc_results[1])) {
            publishDebugInfo("robust_bci_ctrl", "MPC返回NaN结果!", 2);
            return;
        }
        
        // 应用控制命令
        _speed = mpc_results[0];
        _w = mpc_results[1];
        
        // ✅ 发布结果可视化
        ss.str("");
        ss << "MPC结果: speed=" << _speed << ", angular_vel=" << _w;
        publishDebugInfo("robust_bci_ctrl", ss.str(), 0);
        
        // 发布控制命令
        _twist_msg.linear.x = _speed;
        _twist_msg.angular.z = _w;
        _pub_twist.publish(_twist_msg);
        
        // ✅ 发布机器人轨迹
        publishRobotTrajectory(odom.pose.pose.position.x, 
                             odom.pose.pose.position.y, 
                             tf::getYaw(odom.pose.pose.orientation));
        
        // ✅ 发布预测轨迹
        publishPredictionTrajectory(_mpc.mpc_x, _mpc.mpc_y);
        
        publishDebugInfo("robust_bci_ctrl", "MPC控制完成", 0);
        
    } catch (const std::exception& e) {
        // ✅ 异常处理可视化
        std::string error_msg = "MPC控制异常: ";
        error_msg += e.what();
        publishDebugInfo("robust_bci_ctrl", error_msg, 2);
    }

    // const double cmd_v = cmd_user.linear.x;       // Linear velocity
    // const double cmd_w = cmd_user.angular.z;      // angular velocity

   
    // // user's command
    // VectorXd vec_cmd(2);
    // vec_cmd << cmd_v , cmd_w;
  {  //---------原版不知道什么用的代码------------//

    // Update system states: X=[x, y, theta]
    const double px = odom.pose.pose.position.x; //pose: odom frame
    const double py = odom.pose.pose.position.y;
    tf::Pose pose;
    tf::poseMsgToTF(odom.pose.pose, pose);
    const double theta = tf::getYaw(pose.getRotation()); 
    const double v = odom.twist.twist.linear.x; //twist: body fixed frame 
    const double w = odom.twist.twist.angular.z; //angular_velocity

    // // distance and degree 距离障碍物的最近距离和角度
    // const double min_distance = _closest_distance;  
    // const double min_degree = _closest_degree;    

    

        // // state
        // VectorXd state(5);
        // state << 0, 0, 0, min_distance, min_degree;
    _mpc_dif.linear.x = cmd_user.linear.x;
    _mpc_dif.linear.y = cmd_user.angular.z;
        
    _mpc_dif.angular.x = _speed;
    _mpc_dif.angular.y = _w;
    _pub_mpcdif.publish(_mpc_dif);

    std_msgs::Float32 mpc_total_cost;
    mpc_total_cost.data = static_cast<float>(_mpc._mpc_totalcost);
    _pub_totalcost.publish(mpc_total_cost);

    // Display the MPC predicted trajectory
    _mpc_traj = nav_msgs::Path();
    _mpc_traj.header.frame_id = _car_frame; // points in car coordinate        
    _mpc_traj.header.stamp = ros::Time::now();
    for(int i=0; i<_mpc.mpc_x.size(); i++)
    {
        geometry_msgs::PoseStamped tempPose;
        tempPose.header = _mpc_traj.header;
        tempPose.pose.position.x = _mpc.mpc_x[i];
        tempPose.pose.position.y = _mpc.mpc_y[i];
        tempPose.pose.orientation.w = 1.0;
        _mpc_traj.poses.push_back(tempPose); 
    }     
    // publish the mpc trajectory
    _pub_mpctraj.publish(_mpc_traj);

    // cout << "Leader X: " << px << endl;
    // cout << "Leader Y: " << py << endl;
    // cout << "Leader V: " << v << endl;
    // cout << "Leader W: " << w << endl;
    // cout << endl;

}
    }
// CallBack: Update odometry
void RobustMPCNode::odomCB(const nav_msgs::Odometry::ConstPtr& odomMsg)
{
    _odom = *odomMsg;
}

// 读取用户输入的命令 
void RobustMPCNode::read_cmd(const geometry_msgs::Twist::ConstPtr& cmd)
{
    _sub_twist_cmd = *cmd;
}
// 将当前时间清零
void RobustMPCNode::clean_time_CB(const std_msgs::Int32::ConstPtr& msg)
{
    int key = msg->data;
    // cout << "key: " << key << endl;
    if (key == 0) {
        _time_begin = ros::Time::now().toSec();
        while (_time_begin == 0){
            _time_begin = ros::Time::now().toSec();
        }
        // cout << "time_begin： " << _time_begin << endl;
    }

}

// 获取雷达最近距离和角度
void RobustMPCNode::get_laser_callback(const sensor_msgs::LaserScan::ConstPtr& msg)
{
    // ✅ 记录回调开始
    publishDebugInfo("get_laser_callback", "收到激光数据", 0);
    try {
    // 获取距离数组
    std::vector<float> ranges=msg->ranges;
            // ✅ 检查数据有效性
        if (ranges.empty()) {
            publishDebugInfo("get_laser_callback", "激光数据为空!", 1);
            return;
        }
    float angle_min1 = msg -> angle_min;
    float angle_max1 = msg -> angle_max;

    // 获取距离数组的长度
    int len_ = ranges.size();

    // 距离数组标准化
    for(int i=0 ; i < len_ ; i++)
    {
        if(isinf(ranges[i]))
        {
            ranges[i] = 3.5;
        }
        else if(isnan(ranges[i]))
        {
            ranges[i] = 0.0;
        }
    }

    // 求最短距离
    auto closest_range = std::min_element(ranges.begin(),ranges.end());
    float closest_range_ = *(closest_range);

    // 求最短距离所在位置(角度rad)
    auto closest_range_seq = closest_range - ranges.begin() ;
    float closest_degree_rad = angle_min1 + closest_range_seq*(angle_max1-angle_min1)/len_  ;

    // changed
    if((0.785 < closest_degree_rad)&&(closest_degree_rad < 5.495))
    {
        closest_range_ = 3.5;
    }
    //float closest_degree = rad2deg(closest_degree_rad);

    _closest_distance = closest_range_;
    _closest_degree = closest_degree_rad;

    // 显示
    /*
    std::cout << "最短距离：" << closest_range_ << "  ";
    std::cout << "角度：" << closest_degree << "  ";
    std::cout << "弧度：" << closest_degree_rad << std::endl;
    */
    // ✅ 发布处理结果
        std::stringstream ss;
        ss << "最近障碍物: 距离=" << _closest_distance << "m, 角度=" << _closest_degree * 180/M_PI << "度";
        publishDebugInfo("get_laser_callback", ss.str(), 0);
}
catch (const std::exception& e) {
        publishDebugInfo("get_laser_callback", std::string("处理异常: ") + e.what(), 2);
    }
}

{   //--------------------可视化改进-----------------------//
    // 发布调试信息
void RobustMPCNode::publishDebugInfo(const std::string& function_name, const std::string& message, int level) {
    std::string level_str;
    switch(level) {
        case 0: level_str = "INFO"; break;
        case 1: level_str = "WARN"; break;
        case 2: level_str = "ERROR"; break;
        default: level_str = "DEBUG"; break;
    }
    
    std::stringstream ss;
    ss << "[" << getCurrentTimeString() << "][" << level_str << "][" << function_name << "] " << message;
    
    _debug_msg.data = ss.str();
    _pub_debug_info.publish(_debug_msg);
    
    // 同时在控制台输出（可选）
    if (level >= 1) { // 只输出警告和错误
        if (level == 1) ROS_WARN_STREAM(ss.str());
        else if (level == 2) ROS_ERROR_STREAM(ss.str());
    }
}

// 发布MPC状态可视化
void RobustMPCNode::publishMPCStateVisualization(const Eigen::VectorXd& state, const Eigen::VectorXd& cmd) {
    visualization_msgs::Marker marker;
    marker.header.frame_id = _odom_frame;
    marker.header.stamp = ros::Time::now();
    marker.ns = "mpc_state";
    marker.id = _debug_counter++;
    marker.type = visualization_msgs::Marker::TEXT_VIEW_FACING;
    marker.action = visualization_msgs::Marker::ADD;
    
    // 设置位置（机器人当前位置）
    marker.pose.position.x = state[0];
    marker.pose.position.y = state[1];
    marker.pose.position.z = 1.0;
    marker.pose.orientation.w = 1.0;
    
    // 设置文本内容
    std::stringstream text;
    text << "MPC State:\n";
    text << "Pos: (" << std::fixed << std::setprecision(2) << state[0] << ", " << state[1] << ")\n";
    text << "Theta: " << state[2] * 180.0 / M_PI << " deg\n";
    text << "ObsDist: " << state[3] << "m\n";
    text << "ObsAng: " << state[4] * 180.0 / M_PI << " deg\n";
    text << "Cmd: v=" << cmd[0] << ", w=" << cmd[1];
    
    marker.text = text.str();
    marker.scale.z = 0.3;
    marker.color.r = 1.0;
    marker.color.g = 1.0;
    marker.color.b = 0.0;
    marker.color.a = 1.0;
    marker.lifetime = ros::Duration(0.2);
    
    _pub_mpc_state.publish(marker);
}

// 发布障碍物标记
void RobustMPCNode::publishObstacleMarker(double distance, double angle) {
    visualization_msgs::Marker marker;
    marker.header.frame_id = _car_frame;  // 相对于机器人坐标系
    marker.header.stamp = ros::Time::now();
    marker.ns = "obstacle";
    marker.id = 0;
    marker.type = visualization_msgs::Marker::SPHERE;
    marker.action = visualization_msgs::Marker::ADD;
    
    // 计算障碍物在机器人坐标系中的位置
    marker.pose.position.x = distance * cos(angle);
    marker.pose.position.y = distance * sin(angle);
    marker.pose.position.z = 0.0;
    marker.pose.orientation.w = 1.0;
    
    marker.scale.x = 0.5;
    marker.scale.y = 0.5;
    marker.scale.z = 0.5;
    
    // 根据距离改变颜色
    if (distance < 1.0) {
        marker.color.r = 1.0; marker.color.g = 0.0; marker.color.b = 0.0; // 红色：危险
    } else if (distance < 2.0) {
        marker.color.r = 1.0; marker.color.g = 1.0; marker.color.b = 0.0; // 黄色：警告
    } else {
        marker.color.r = 0.0; marker.color.g = 1.0; marker.color.b = 0.0; // 绿色：安全
    }
    marker.color.a = 0.8;
    marker.lifetime = ros::Duration(0.5);
    
    _pub_obstacle_marker.publish(marker);
}

// 发布机器人轨迹
void RobustMPCNode::publishRobotTrajectory(double x, double y, double theta) {
    // 限制轨迹长度（避免内存泄漏）
    if (_robot_path.poses.size() > 1000) {
        _robot_path.poses.erase(_robot_path.poses.begin());
    }
    
    geometry_msgs::PoseStamped pose;
    pose.header.frame_id = _odom_frame;
    pose.header.stamp = ros::Time::now();
    pose.pose.position.x = x;
    pose.pose.position.y = y;
    
    tf2::Quaternion q;
    q.setRPY(0, 0, theta);
    pose.pose.orientation = tf2::toMsg(q);
    
    _robot_path.poses.push_back(pose);
    _robot_path.header.frame_id = _odom_frame;
    _robot_path.header.stamp = ros::Time::now();
    
    _pub_robot_trajectory.publish(_robot_path);
}}

/*****************/
/* MAIN FUNCTION */
/*****************/
int main(int argc, char **argv)
{
    // ✅ 步骤1：ROS初始化
    ros::init(argc, argv, "RobustleaderMPCNode");
    // ✅ 步骤2：创建对象（执行构造函数）
    RobustMPCNode mpc_node;
    // ✅ 步骤3：启动异步spinner
    ros::AsyncSpinner spinner(mpc_node.get_thread_numbers()); // Use multi threads
    spinner.start();
    // 🔥 现在ROS开始在后台监听消息和定时器
    
    // ✅ 步骤4：主线程等待
    ros::waitForShutdown();
    return 0;
}
